#!/bin/env python
import sys

def gen_chroma_singularity_define_files(chroma_src_path = "/home/guilongcheng/lattice_cloud/source_codes",
                                        cuda_imag="update_cuda11.8_ubuntu22.04",
                                        sm="sm_35"):
    chroma_addon_def=f"""Bootstrap:localimage
From: {cuda_imag}.sif

%setup
    mkdir -p /install/
    mkdir -p /src
%files
    {chroma_src_path}/* /src/
%post
    cd /src/
    ./build.sh {sm} 

    cd /src/chroma_addons
    sh install.sh {sm} qdpjit 
    mkdir -p /install/chroma_addon/bin
    cp /build/chroma_addon_{sm}/chroma_addon /install/chroma_addon/bin
%environment
    export PATH=/install/chroma-jit-{sm}/bin:$PATH
    export PATH=/install/chroma_addon/bin:$PATH
%runscript
    """
    chroma_addon_imag = "chroma_addon_qdpjit_quda_"+ sm + "_" + cuda_imag
    with open(chroma_addon_imag + ".def","w") as f:
        f.writelines(chroma_addon_def)

for i in ["sm_35","sm_60","sm_70","sm_80"]:
    gen_chroma_singularity_define_files(sm=i)    
